#include <userver/congestion_control/component.hpp>
#include <userver/congestion_control/controller.hpp>

#include <congestion_control/watchdog.hpp>
#include <userver/server/congestion_control/sensor.hpp>

#include <userver/components/component.hpp>
#include <userver/components/statistics_storage.hpp>
#include <userver/concurrent/async_event_channel.hpp>
#include <userver/dynamic_config/storage/component.hpp>
#include <userver/dynamic_config/value.hpp>
#include <userver/hostinfo/cpu_limit.hpp>
#include <userver/logging/log.hpp>
#include <userver/server/component.hpp>
#include <userver/server/server.hpp>
#include <userver/yaml_config/merge_schemas.hpp>

#include <dynamic_config/variables/USERVER_RPS_CCONTROL_ACTIVATED_FACTOR_METRIC.hpp>
#include <dynamic_config/variables/USERVER_RPS_CCONTROL_ENABLED.hpp>

#ifndef ARCADIA_ROOT
#include "generated/src/congestion_control/component.yaml.hpp"  // Y_IGNORE
#endif

USERVER_NAMESPACE_BEGIN

namespace congestion_control {

namespace {

const auto kServerControllerName = "server-main-tp-cc";

void FormatStats(const Controller& c, size_t activated_factor, utils::statistics::Writer& writer) {
    writer["is-enabled"] = c.IsEnabled() ? 1 : 0;

    auto limit = c.GetLimit();
    writer["is-activated"] = (limit.load_limit && limit.current_load * activated_factor < *limit.load_limit) ? 1 : 0;
    if (limit.load_limit) {
        writer["limit"] = *limit.load_limit;
    }

    const auto& stats = c.GetStats();
    {
        auto builder_states = writer["states"];
        builder_states["no-limit"] = stats.no_limit;
        builder_states["not-overloaded-no-pressure"] = stats.not_overload_no_pressure;
        builder_states["not-overloaded-under-pressure"] = stats.not_overload_pressure;
        builder_states["overloaded-no-pressure"] = stats.overload_no_pressure;
        builder_states["overloaded-under-pressure"] = stats.overload_pressure;
    }

    auto diff = std::chrono::steady_clock::now().time_since_epoch() - stats.last_overload_pressure.load();
    writer["time-from-last-overloaded-under-pressure-secs"] =
        std::chrono::duration_cast<std::chrono::seconds>(diff).count();
    writer["current-state"] = stats.current_state;
}

}  // namespace

struct Component::Impl {
    dynamic_config::Source dynamic_config;
    server::Server& server;
    server::congestion_control::Sensor server_sensor;
    server::congestion_control::Limiter server_limiter;
    Controller server_controller;

    std::atomic<bool> fake_mode;
    std::atomic<bool> force_disabled{false};
    std::atomic<size_t> last_activate_factor{1};

    // These subscriptions and tasks must be the last fields!
    Watchdog wd;
    utils::statistics::Entry statistics_holder;
    concurrent::AsyncEventSubscriberScope config_subscription;
    // See the comment above before adding new fields.

    Impl(dynamic_config::Source dynamic_config, server::Server& server, engine::TaskProcessor& tp, bool fake_mode)
        : dynamic_config(dynamic_config),
          server(server),
          server_sensor(tp),
          server_controller(kServerControllerName, dynamic_config),
          fake_mode(fake_mode)
    {
        server_limiter.RegisterLimitee(server);
        server_sensor.RegisterRequestsSource(server);
    }
};

Component::Component(const components::ComponentConfig& config, const components::ComponentContext& context)
    : components::ComponentBase(config, context),
      pimpl_(
          context.FindComponent<components::DynamicConfig>().GetSource(),
          context.FindComponent<components::Server>().GetServer(),
          engine::current_task::GetTaskProcessor(),
          config["fake-mode"].As<bool>(false)
      )
{
    auto min_threads = config["min-cpu"].As<size_t>(1);
    auto only_rtc = config["only-rtc"].As<bool>(true);

    pimpl_->server.SetRpsRatelimitStatusCode(static_cast<server::http::HttpStatus>(config["status-code"].As<int>(429)));

    if (!pimpl_->fake_mode && only_rtc && !hostinfo::IsInRtc()) {
        throw std::runtime_error(
            "Started outside of RTC, don't know how to get number of CPU cores. You have to pass CPU_LIMIT environment "
            "variable with the correct value in order CC to work properly."
        );
    }

    auto cpu_limit = hostinfo::CpuLimit().value_or(0);
    if (!pimpl_->fake_mode && cpu_limit < min_threads) {
        LOG_WARNING() << fmt::format(
            "current cpu limit ({}) less that minimum limit ({}) "
            "for Congestion Control, forcing fake-mode",
            cpu_limit,
            min_threads
        );
        pimpl_->fake_mode = true;
    }

    if (pimpl_->fake_mode) {
        LOG_WARNING()
            << "congestion_control is started in fake-mode, no RPS limit "
               "is enforced";
    }

    pimpl_->wd.Register({pimpl_->server_sensor, pimpl_->server_limiter, pimpl_->server_controller});

    pimpl_->config_subscription = pimpl_->dynamic_config.UpdateAndListen(this, kName, &Component::OnConfigUpdate);

    auto& storage = context.FindComponent<components::StatisticsStorage>().GetStorage();
    pimpl_->statistics_holder = storage.RegisterWriter(std::string{kName}, [this](utils::statistics::Writer& writer) {
        ExtendWriter(writer);
    });
}

Component::~Component() {
    pimpl_->config_subscription.Unsubscribe();
    pimpl_->statistics_holder.Unregister();
}

void Component::OnConfigUpdate(const dynamic_config::Snapshot& cfg) {
    const bool is_enabled_dynamic = cfg[::dynamic_config::USERVER_RPS_CCONTROL_ENABLED];
    pimpl_->last_activate_factor = cfg[::dynamic_config::USERVER_RPS_CCONTROL_ACTIVATED_FACTOR_METRIC];

    bool enabled = !pimpl_->fake_mode.load() && !pimpl_->force_disabled.load();
    if (enabled && !is_enabled_dynamic) {
        LOG_INFO()
            << "Congestion control is explicitly disabled in "
               "USERVER_RPS_CCONTROL_ENABLED config";
        enabled = false;
    }
    pimpl_->server_controller.SetEnabled(enabled);
}

void Component::OnAllComponentsLoaded() {
    LOG_DEBUG()
        << "Found " << pimpl_->server.GetThrottlableHandlersCount()
        << " registered HTTP handlers with enabled throttling";
    if (pimpl_->server.GetThrottlableHandlersCount() == 0) {
        pimpl_->force_disabled = true;
        LOG_WARNING() << "No throttlable HTTP handlers registered, disabling";

        // apply fake_mode
        OnConfigUpdate(pimpl_->dynamic_config.GetSnapshot());
    }
}

void Component::OnAllComponentsAreStopping() { pimpl_->wd.Stop(); }

void Component::ExtendWriter(utils::statistics::Writer& writer) {
    if (!pimpl_->force_disabled) {
        auto rps = writer["rps"];
        FormatStats(pimpl_->server_controller, pimpl_->last_activate_factor, rps);
    }
}

server::congestion_control::Limiter& Component::GetServerLimiter() { return pimpl_->server_limiter; }

server::congestion_control::Sensor& Component::GetServerSensor() { return pimpl_->server_sensor; }

const congestion_control::Controller& Component::GetServerController() const { return pimpl_->server_controller; }

yaml_config::Schema Component::GetStaticConfigSchema() {
    return yaml_config::MergeSchemasFromResource<components::ComponentBase>("src/congestion_control/component.yaml");
}

}  // namespace congestion_control

USERVER_NAMESPACE_END
